/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2021-2022. All rights reserved.
 * Description: HDDS接口
 * Author: liutianming
 * Create: 2021-11-15
 */

#include <sys/types.h>
#include <unistd.h>
#include <mutex>
#include <memory>
#include "log.h"
#include "hccl/base.h"
#include "config.h"
#include "hdcs_impl.h"
#include "driver/ascend_hal.h"
#include "hccl_ex.h"

using namespace std;
using namespace hccl;

static std::unordered_map<string, HdcsImpl*> g_groupMapImpl{};
static bool g_enableEntryLog{ false };

HcclResult DevGetPid(int *pid)
{
    *pid = getpid();
    HCCL_RUN_INFO("[Hdcs][GetPid]hdcs get pid %d", *pid);
    return HCCL_SUCCESS;
}

HcclResult DevCommInit(HdcsCsInitPara para)
{
    HCCL_DEBUG("call DevCommInit");
    HCCL_DEBUG("psSize[%u]", para.psSize);
    for (u32 i = 0; i < para.psSize; ++i) {
        HCCL_INFO("psList: index[%u]=[%u]", i, para.psList[i]);
    }
    HCCL_DEBUG("rank[%u]", para.rank);
    HCCL_DEBUG("para.rankTableM[%s]", para.rankTableM);
    HCCL_DEBUG("para.groupName[%s]", para.groupName);

    g_enableEntryLog = para.enableEntryLog;
    HCCL_RUN_INFO("[Dev][CommInit]enableEntryLog[%u]", para.enableEntryLog);
    std::string groupName = para.groupName;
    if (g_groupMapImpl.find(groupName) == g_groupMapImpl.end()) {
        std::vector<u32> psRankId(para.psSize);
        for (u32 index = 0; index < para.psSize; index++) {
            // psId 2 rankId
            psRankId[index] = para.psList[index];
        }

        HdcsImpl* impl = new (std::nothrow) HdcsImpl();
        CHK_PTR_NULL(impl);
        HcclResult ret = impl->Init(para, psRankId);
        if (ret != HCCL_SUCCESS) {
            HCCL_ERROR("HdcsImpl init fail, ret is %d", ret);
            delete impl;
            return ret;
        }
        g_groupMapImpl[groupName] = impl;
    }

    HCCL_RUN_INFO("[Dev][CommInit]hdcs Init success");
    return HCCL_SUCCESS;
}

HcclResult DevCommDeInit()
{
    HCCL_RUN_INFO("[Dev][CommDeInit]hdcs DeInit Entry");
    for (auto iter : g_groupMapImpl) {
        if (iter.second != nullptr) {
            delete iter.second;
        }
    }

    g_groupMapImpl.clear();
    return HCCL_SUCCESS;
}

HcclResult RegTransport(HdcsRegTransportPara para)
{
    HCCL_RUN_INFO("[Reg][Transport]hdcs RegTransport Entry");
    std::string groupName = para.groupName;

    if (g_groupMapImpl.find(groupName) == g_groupMapImpl.end()) {
        HCCL_INFO("[Reg][Transport]hdcs not Init !!!");
    }

    HdcsImpl* impl = g_groupMapImpl[groupName];
    impl->RegTransport(para);
    return HCCL_SUCCESS;
}

HcclResult IsendLookupRequest(void *embeddingData)
{
    CHK_PTR_NULL(embeddingData);
    HdcsEmbeddingServiceParam *embeddingInfo = reinterpret_cast<HdcsEmbeddingServiceParam *>(embeddingData);

    HCCL_RUN_INFO("[Isend][LookupRequest]hdcs RemoteLookUp Entry");
    std::string groupName = embeddingInfo->groupName;
    if (g_groupMapImpl.find(groupName) == g_groupMapImpl.end()) {
        HCCL_ERROR("[Isend][LookupRequest]hdcs not Init !!!");
        return HCCL_E_PARA;
    }

    HdcsImpl* impl = g_groupMapImpl[groupName];

    CHK_RET(impl->RemoteLookUp(embeddingInfo));
    HCCL_RUN_INFO("hdcs RemoteLookUp Entry succ");
    return HCCL_SUCCESS;
}

HcclResult IrecvLookupResponse(HcclRequest* requestArray)
{
    CHK_PTR_NULL(requestArray);
    return HCCL_SUCCESS;
}

HcclResult IsendUpdateRequest(void *embeddingData)
{
    CHK_PTR_NULL(embeddingData);
    HdcsEmbeddingServiceParam *embeddingInfo = reinterpret_cast<HdcsEmbeddingServiceParam *>(embeddingData);

    HCCL_RUN_INFO("[Isend][UpdateRequest]hdcs RemoteUpdate Entry");
    std::string groupName = embeddingInfo->groupName;
    if (g_groupMapImpl.find(groupName) == g_groupMapImpl.end()) {
        HCCL_ERROR("[Isend][UpdateRequest]hdcs not Init !!!");
        return HCCL_E_PARA;
    }

    HdcsImpl* impl = g_groupMapImpl[groupName];

    CHK_RET(impl->RemoteUpdate(embeddingInfo));
    HCCL_RUN_INFO("[Isend][UpdateRequest]hdcs RemoteUpdate Entry succ");

    return HCCL_SUCCESS;
}

HcclResult IrecvUpdateResponse(HcclRequest* request)
{
    CHK_PTR_NULL(request);
    return HCCL_SUCCESS;
}

HcclResult TestSome(HcclRequest* request)
{
    CHK_PTR_NULL(request);
    return HCCL_SUCCESS;
}

HcclResult WaitSome()
{
    return HCCL_SUCCESS;
}

static u32 CollGetTableId(void *tableIdAddr, bool zeroCpyPtr)
{
    if (UNLIKELY(tableIdAddr == nullptr)) {
        return ~0U;
    }

    if (!zeroCpyPtr) {
        return *(static_cast<u32 *>(tableIdAddr));
    }

    u32 *tableId = reinterpret_cast<u32 *>(*(static_cast<u64 *>(tableIdAddr)));
    if (UNLIKELY(tableId == nullptr)) {
        return ~0U;
    } else {
        return *tableId;
    }
}

HcclResult CollLookupKeysDuplicates(void *pipelineParams)
{
    CHK_PTR_NULL(pipelineParams);
    HdcsEmbeddingServiceParam *lookupInfo = reinterpret_cast<HdcsEmbeddingServiceParam *>(pipelineParams);

    u32 tableId = CollGetTableId(lookupInfo->tableIdAddr, lookupInfo->intZerocpyFlag);
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-CollLookupKeysDuplicates tag[%d], tableId[%u], cubeIndex[%d] pipelineKeyNum[%llu]",
            lookupInfo->tag, tableId, lookupInfo->cubeIndex, lookupInfo->pipelineKeyNum);
    }

    std::string groupName = lookupInfo->groupName;
    auto impl = g_groupMapImpl.find(groupName);
    if (impl == g_groupMapImpl.end()) {
        HCCL_ERROR("[Isend][UpdateRequest]hdcs not Init !!!");
        return HCCL_E_PARA;
    }

    TIME_PRINT(CHK_RET(impl->second->CollLookupKeysDuplicates(lookupInfo)));
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-CollLookupKeysDuplicates succ, tag[%d], tableId[%u]", lookupInfo->tag, tableId);
    }

    return HCCL_SUCCESS;
}

HcclResult CollLookupSendKeys(void *pipelineParams)
{
    CHK_PTR_NULL(pipelineParams);
    HdcsLookupPipelineParams *lookupInfo = reinterpret_cast<HdcsLookupPipelineParams *>(pipelineParams);

    u32 tableId = CollGetTableId(lookupInfo->tableIdAddr, lookupInfo->intZerocpyFlag);
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-CollLookupSendKeys tag[%d], tableId[%u], cubeIndex[%d]", lookupInfo->tag, tableId,
            lookupInfo->cubeIndex);
    }

    std::string groupName = lookupInfo->groupName;
    if (g_groupMapImpl.find(groupName) == g_groupMapImpl.end()) {
        HCCL_ERROR("[Isend][UpdateRequest]hdcs not Init !!!");
        return HCCL_E_PARA;
    }

    HdcsImpl* impl = g_groupMapImpl[groupName];
    TIME_PRINT(CHK_RET(impl->CollLookupSendKeys(lookupInfo)));
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-CollLookupSendKeys succ, tag[%d], tableId[%u]", lookupInfo->tag, tableId);
    }

    return HCCL_SUCCESS;
}

HcclResult CollLookupRecvValues(void *pipelineParams)
{
    CHK_PTR_NULL(pipelineParams);
    HdcsLookupPipelineParams *lookupInfo = reinterpret_cast<HdcsLookupPipelineParams *>(pipelineParams);

    u32 tableId = CollGetTableId(lookupInfo->tableIdAddr, lookupInfo->intZerocpyFlag);
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-CollLookupRecvValues tag[%d], tableId[%u], cubeIndex[%d]", lookupInfo->tag,
            tableId, lookupInfo->cubeIndex);
    }

    std::string groupName = lookupInfo->groupName;
    if (g_groupMapImpl.find(groupName) == g_groupMapImpl.end()) {
        HCCL_ERROR("[Isend][UpdateRequest]hdcs not Init !!!");
        return HCCL_E_PARA;
    }

    HdcsImpl* impl = g_groupMapImpl[groupName];
    TIME_PRINT(CHK_RET(impl->CollLookupRecvValues(lookupInfo)));
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-CollLookupRecvValues succ, tag[%d], tableId[%u]", lookupInfo->tag, tableId);
    }

    return HCCL_SUCCESS;
}

HcclResult CollLookupRecoverValue(void *pipelineParams)
{
    CHK_PTR_NULL(pipelineParams);
    HdcsLookupPipelineParams *lookupInfo = reinterpret_cast<HdcsLookupPipelineParams *>(pipelineParams);
    u32 tableId = CollGetTableId(lookupInfo->tableIdAddr, lookupInfo->intZerocpyFlag);
    HCCL_RUN_INFO("Entry-CollLookupRecoverValue tag[%d], tableId[%u], cubeIndex[%d]", lookupInfo->tag, tableId,
        lookupInfo->cubeIndex);

    std::string groupName = lookupInfo->groupName;
    if (g_groupMapImpl.find(groupName) == g_groupMapImpl.end()) {
        HCCL_ERROR("[Isend][UpdateRequest]hdcs not Init !!!");
        return HCCL_E_PARA;
    }

    HdcsImpl* impl = g_groupMapImpl[groupName];
    TIME_PRINT(CHK_RET(impl->CollLookupRecoverValue(lookupInfo)));
    HCCL_RUN_INFO("Entry-CollLookupRecoverValue succ, tag[%d], tableId[%u]", lookupInfo->tag, tableId);

    return HCCL_SUCCESS;
}

HcclResult CollLookupResetUniqueHandle(void *pipelineParams)
{
    CHK_PTR_NULL(pipelineParams);
    HdcsLookupPipelineParams *lookupInfo = reinterpret_cast<HdcsLookupPipelineParams *>(pipelineParams);
    u32 tableId = CollGetTableId(lookupInfo->tableIdAddr, lookupInfo->intZerocpyFlag);
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-CollLookupResetUniqueHandle tag[%d], tableId[%u], cubeIndex[%d]",
            lookupInfo->tag, tableId, lookupInfo->cubeIndex);
    }

    std::string groupName = lookupInfo->groupName;
    if (g_groupMapImpl.find(groupName) == g_groupMapImpl.end()) {
        HCCL_ERROR("[Isend][UpdateRequest]hdcs not Init !!!");
        return HCCL_E_PARA;
    }

    HdcsImpl* impl = g_groupMapImpl[groupName];
    TIME_PRINT(CHK_RET(impl->CollLookupResetUniquesHandle(lookupInfo)));
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-CollLookupResetUniqueHandle succ, tag[%d], tableId[%u]", lookupInfo->tag, tableId);
    }

    return HCCL_SUCCESS;
}

HcclResult CollLookupWaitSendKeyFinish(void *pipelineParams)
{
    CHK_PTR_NULL(pipelineParams);
    HdcsLookupPipelineParams *lookupInfo = reinterpret_cast<HdcsLookupPipelineParams *>(pipelineParams);
    u32 tableId = CollGetTableId(lookupInfo->tableIdAddr, lookupInfo->intZerocpyFlag);
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-CollLookupWaitSendKeyFinish tag[%d], tableId[%u], cubeIndex[%d]",
            lookupInfo->tag, tableId, lookupInfo->cubeIndex);
    }

    std::string groupName = lookupInfo->groupName;
    if (g_groupMapImpl.find(groupName) == g_groupMapImpl.end()) {
        HCCL_ERROR("[Isend][UpdateRequest]hdcs not Init !!!");
        return HCCL_E_PARA;
    }

    HdcsImpl* impl = g_groupMapImpl[groupName];
    TIME_PRINT(CHK_RET(impl->CollLookupWaitSendKeyFinish(lookupInfo)));
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-CollLookupWaitSendKeyFinish succ, tag[%d], tableId[%u]", lookupInfo->tag, tableId);
    }

    return HCCL_SUCCESS;
}

HcclResult CollLookupGatherFinish(void *pipelineParams)
{
    CHK_PTR_NULL(pipelineParams);
    HdcsLookupPipelineParams *lookupInfo = reinterpret_cast<HdcsLookupPipelineParams *>(pipelineParams);
    u32 tableId = CollGetTableId(lookupInfo->tableIdAddr, lookupInfo->intZerocpyFlag);
    HCCL_RUN_INFO("Entry-CollLookupGatherFinish succ tag[%d], tableId[%u]", lookupInfo->tag, tableId);

    return HCCL_SUCCESS;
}

HcclResult RemoteUpdateKeyReduce(void *embeddingData)
{
    CHK_PTR_NULL(embeddingData);
    HdcsEmbeddingServiceParam *embeddingInfo = reinterpret_cast<HdcsEmbeddingServiceParam *>(embeddingData);

    u32 tableId = CollGetTableId(embeddingInfo->tableIdAddr, embeddingInfo->intZerocpyFlag);
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-RemoteUpdateKeyReduce, tag[%d], tableId[%u]",
            embeddingInfo->tag, tableId);
    }
    std::string groupName = embeddingInfo->groupName;
    if (g_groupMapImpl.find(groupName) == g_groupMapImpl.end()) {
        HCCL_ERROR("[Isend][UpdateRequest]hdcs not Init !!!");
        return HCCL_E_PARA;
    }

    HdcsImpl* impl = g_groupMapImpl[groupName];
    TIME_PRINT(CHK_RET(impl->RemoteUpdateKeyReduce(embeddingInfo)));
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-RemoteUpdateKeyReduce succ, tag[%d], tableId[%u]",
            embeddingInfo->tag, tableId);
    }

    return HCCL_SUCCESS;
}

HcclResult RemoteUpdateReduceSum(void *embeddingData)
{
    CHK_PTR_NULL(embeddingData);
    HdcsEmbeddingServiceParam *embeddingInfo = reinterpret_cast<HdcsEmbeddingServiceParam *>(embeddingData);

    u32 tableId = CollGetTableId(embeddingInfo->tableIdAddr, embeddingInfo->intZerocpyFlag);
    HCCL_RUN_INFO("Entry-RemoteUpdateReduceSum, tag[%d], tableId[%u]",
        embeddingInfo->tag, tableId);
    std::string groupName = embeddingInfo->groupName;
    if (g_groupMapImpl.find(groupName) == g_groupMapImpl.end()) {
        HCCL_ERROR("[Isend][UpdateRequest]hdcs not Init !!!");
        return HCCL_E_PARA;
    }

    HdcsImpl* impl = g_groupMapImpl[groupName];
    TIME_PRINT(CHK_RET(impl->RemoteUpdateReduceSum(embeddingInfo)));
    HCCL_RUN_INFO("Entry-RemoteUpdateReduceSum succ, tag[%d], tableId[%u]",
        embeddingInfo->tag, tableId);

    return HCCL_SUCCESS;
}

HcclResult RemoteUpdateSendRequest(void *embeddingData)
{
    CHK_PTR_NULL(embeddingData);
    HdcsEmbeddingServiceParam *embeddingInfo = reinterpret_cast<HdcsEmbeddingServiceParam *>(embeddingData);

    u32 tableId = CollGetTableId(embeddingInfo->tableIdAddr, embeddingInfo->intZerocpyFlag);

    s64 globalStep = (embeddingInfo->globalStepAddr == nullptr) ? DEFAULT_GLOBAL_STEP_VALUE :\
        *(static_cast<s64 *>(embeddingInfo->globalStepAddr));
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-RemoteUpdateSendRequest, tag[%d], tableId[%u] globalStep[%lld]",
            embeddingInfo->tag, tableId, globalStep);
    }

    std::string groupName = embeddingInfo->groupName;
    if (g_groupMapImpl.find(groupName) == g_groupMapImpl.end()) {
        HCCL_ERROR("[Isend][UpdateRequest]hdcs not Init !!!");
        return HCCL_E_PARA;
    }

    HdcsImpl* impl = g_groupMapImpl[groupName];
    TIME_PRINT(CHK_RET(impl->RemoteUpdateSendRequest(embeddingInfo)));
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-RemoteUpdateSendRequest succ, tag[%d], tableId[%u]",
            embeddingInfo->tag, tableId);
    }

    return HCCL_SUCCESS;
}
HcclResult RemoteUpdateResetUniqueHandle(void *embeddingData)
{
    CHK_PTR_NULL(embeddingData);
    HdcsEmbeddingServiceParam *embeddingInfo = reinterpret_cast<HdcsEmbeddingServiceParam *>(embeddingData);

    u32 tableId = CollGetTableId(embeddingInfo->tableIdAddr, embeddingInfo->intZerocpyFlag);
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-RemoteUpdateResetUniqueHandle, tag[%d], tableId[%u]",
            embeddingInfo->tag, tableId);
    }

    std::string groupName = embeddingInfo->groupName;
    if (g_groupMapImpl.find(groupName) == g_groupMapImpl.end()) {
        HCCL_ERROR("[Isend][UpdateRequest]hdcs not Init !!!");
        return HCCL_E_PARA;
    }

    HdcsImpl* impl = g_groupMapImpl[groupName];
    TIME_PRINT(CHK_RET(impl->RemoteUpdateResetUniqueHandle(embeddingInfo)));
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-RemoteUpdateResetUniqueHandle succ, tag[%d], tableId[%u]",
            embeddingInfo->tag, tableId);
    }

    return HCCL_SUCCESS;
}

HcclResult RemoteUpdateRecvResponse(void *embeddingData)
{
    CHK_PTR_NULL(embeddingData);
    HdcsEmbeddingServiceParam *embeddingInfo = reinterpret_cast<HdcsEmbeddingServiceParam *>(embeddingData);

    u32 tableId = CollGetTableId(embeddingInfo->tableIdAddr, embeddingInfo->intZerocpyFlag);
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-RemoteUpdateRecvResponse, tag[%d], tableId[%u]",
            embeddingInfo->tag, tableId);
    }

    std::string groupName = embeddingInfo->groupName;
    if (g_groupMapImpl.find(groupName) == g_groupMapImpl.end()) {
        HCCL_ERROR("[Isend][UpdateRequest]hdcs not Init !!!");
        return HCCL_E_PARA;
    }

    HdcsImpl* impl = g_groupMapImpl[groupName];
    TIME_PRINT(CHK_RET(impl->RemoteUpdateRecvResponse(embeddingInfo)));
    if (g_enableEntryLog) {
        HCCL_RUN_INFO("Entry-RemoteUpdateRecvResponse succ, tag[%d], tableId[%u]",
            embeddingInfo->tag, tableId);
    }

    return HCCL_SUCCESS;
}

HcclResult ChkServiceCancel(void *interfaceParams)
{
    CHK_PTR_NULL(interfaceParams);
    HdcsEmbeddingServiceParam *embeddingInfo = reinterpret_cast<HdcsEmbeddingServiceParam *>(interfaceParams);
    HCCL_RUN_INFO("[ChkServiceCancel]hdcs ChkServiceCancel Entry errorMsg[%p]", embeddingInfo->errorMsg);
    std::string groupName = embeddingInfo->groupName;
    if (g_groupMapImpl.find(groupName) == g_groupMapImpl.end()) {
        HdcsCsInitPara para;
        para.psSize = 0;
        para.workerSize = 0;
        para.rank = 0;
        para.isInference = true;
        CHK_SAFETY_FUNC_RET(strcpy_s(para.groupName, GROUP_NAME_MAX_LEN, groupName.c_str()));
        DevCommInit(para);
    }

    HdcsImpl* impl = g_groupMapImpl[groupName];
    CHK_PTR_NULL(impl);
    TIME_PRINT(CHK_RET(impl->ChkServiceCancel(embeddingInfo)));
    HCCL_RUN_INFO("[ChkServiceCancel]hdcs ChkServiceCancel Entry succ");
    return HCCL_SUCCESS;
}